import numarray as N
import numarray.mlab as mlab
import numarray.nd_image as nd
import numarray.convolve as conv
import glob,pyfits,os,sys

#$Id: undopuft.py 15 2006-04-11 23:22:57Z magee $
# -------------------------------------------------------------
# This program attempts to remove the "Mr. Staypuft" anomaly which occurs with
# NICMOS observations.  The approximate effect of this anomaly is to create 
# faint ghost images at 128 pixel intervals in the other chips.  This is thought
# to occur due to the interplay between the readout processes (and the voltage of
# the power supply).
#
# translated from original program undopuft.pro from the Eddie Bergeron
# and runpuft.pro written by Susan Kassin under supervision of Eddie Bergeron
#
# translation done by Rychard Bouwens 12/14/2005
#

__version__      = '$Revision: 15 $ '[11:-3]
__version_date__ = '$Date: 2006-04-11 23:22:57 +0000 (Tue, 11 Apr 2006) $ '[7:-3]
__author__       = 'R. Bouwens, <bouwens@ucolick.org>, D. Magee, <magee@ucolick.org>'


def undopuftdriver(im,cam):
  # camera specifies the NIC camera

  cc=N.zeros((256,256),type=N.Float32)

  if cam==1:
    im=mlab.rot90(im,1)
    hh=0.8  # half_life, in pix
    bmax=87. # max puft signal
    tx=2   # number of rows over which sig decays
    ampscale= (1.30024e-05) *20. # sig to amp scale factor
    a1=(-0.00000) # quadratic slope
    a2=0.003  # quadratic zeropoint
    la=0 # quadratic look-ahead (in pixels)

  if cam==2:
    hh=0.8  # half_life, in pix
    bmax=87. # max puft signal
    tx=2   # number of rows over which sig decays
    ampscale= (1.30024e-05)*20. # sig to amp scale factor
    a1=(-0.00014) # quadratic slope
    a2=0.07 # quadratic zeropoint
    la=128 # quadratic look-ahead (in pixels)

  if cam==3:
    im=mlab.rot90(im,2)
    hh=0.8  # half_life, in pix
    bmax=87. # max puft signal
    tx=2   # number of rows over which sig decays
    ampscale= (1.30024e-05)*20. # sig to amp scale factor
    a1=(-0.00014) # quadratic slope
    a2=abs((128*tx*a1))*1.0  # quadratic zeropoint
    la=128 # quadratic look-ahead (in pixels)

  # ********************************************************************

  print 'unwrapping quads'

  q1 = N.transpose(im[0:128,0:128]).flat[::-1]
  q2 = N.transpose(im[0:128,128:256]).flat[::-1]
  q3 = N.transpose(im[128:256,128:256]).flat[::-1]
  q4 = N.transpose(im[128:256,0:128]).flat[::-1]

  totsig = abs((q1+q2+q3+q4)*ampscale)
  totsigla=totsig[la:] # look ahead
  # but can't look ahead for the last row, so ju st use the original value and hope for the best
  totsigla=N.concatenate((totsigla,totsig[16384-la:]))

#  b=N.zeros((16384,),type=N.Float32)

  print 'calculating correction'

  fv0 = N.exp(-N.arange(128*tx)/hh)
  fv1 = N.arange(128.*tx)*a1+a2
  nsigend = tx*128
  readnoiselimit = 40.0 * ampscale
  # jump over readnoise
  mask = (totsig > readnoiselimit)
  totsig *= mask
  totsigla *= mask
  b = conv.convolve(totsig,fv0)[:16384] + conv.convolve(totsigla,fv1)[:16384]
  b = N.clip(b,0,bmax)

  print 'reforming result'
  # reform into a quad
  cc[0:128,0:128] = N.transpose(N.reshape(b[::-1],(128,128)))
  cc[0:128,128:] = cc[0:128,0:128]
  cc[128:,0:128] = cc[0:128,0:128]
  cc[128:,128:] = cc[0:128,0:128]

  # unrotate cams if necessary *****************************************

  if cam == 1:
    cc = mlab.rot90(cc,3)
  if cam == 3:
    cc = mlab.rot90(cc,2)

  # ********************************************************************

  return cc

def undopuft(fin, fout, dark):
  file = pyfits.open(fin)
  ncals = (len(file) - 1) / 5
  darkfile = dark
  print 'working on ', fin,', darkfile: ',darkfile
  camera = file[0].header['CAMERA']
  for i in range(ncals):
    ii = file['SCI',i+1].data.copy()
    ii -= file['SCI', ncals].data
    dark = pyfits.getdata(darkfile, ext='SCI', extver=i+1)
    ccall = undopuftdriver(ii-dark, camera)
    file['SCI', i+1].data = file['SCI', i+1].data.astype(N.Float32)
    file['SCI', i+1].data -= ccall
  print 'writing ', fout
  file.writeto(fout)
  file.close()

if __name__ == '__main__':
  raw = sys.argv[1]
  dark = sys.argv[2]
  raw_out = raw.replace('raw0.fits', 'raw.fits')
  undopuft(raw, raw_out, dark)
